package com.jvmmonitor.jmx;

import java.lang.management.*;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.CopyOnWriteArrayList;
import java.util.concurrent.locks.ReentrantReadWriteLock;

import javax.management.*;
import javax.management.openmbean.CompositeData;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.jvmmonitor.model.GcStats;
import com.jvmmonitor.model.MemoryStats;
import com.jvmmonitor.service.ObjectSizeMonitor;

/**
 * JMX监控管理器
 */
public class JmxMonitor {
    private static final Logger logger = LoggerFactory.getLogger(JmxMonitor.class);
    
    // JMX MBean服务器
    private MBeanServer mbeanServer;
    
    // 内存MXBean
    private MemoryMXBean memoryMXBean;
    
    // 内存池MXBean列表
    private List<MemoryPoolMXBean> memoryPoolMXBeans;
    
    // GC MXBean列表
    private List<GarbageCollectorMXBean> garbageCollectorMXBeans;
    
    // 最新的内存统计数据
    private volatile MemoryStats latestMemoryStats;
    
    // 历史内存统计数据 - 使用CopyOnWriteArrayList提高并发性能
    private List<MemoryStats> memoryStatsHistory;
    
    // GC统计数据 - 使用CopyOnWriteArrayList提高并发性能
    private List<GcStats> gcStatsHistory;
    
    // 历史数据最大数量
    private static final int MAX_HISTORY_SIZE = 100;
    
    // 读写锁，用于优化历史数据访问
    private final ReentrantReadWriteLock memoryStatsLock = new ReentrantReadWriteLock();
    private final ReentrantReadWriteLock gcStatsLock = new ReentrantReadWriteLock();
    
    // GC通知监听器
    private NotificationListener gcNotificationListener;
    
    // 通知过滤器
    private NotificationFilter gcNotificationFilter;
    
    // 年轻代内存池
    private MemoryPoolMXBean youngMemoryPool;
    private MemoryPoolMXBean oldMemoryPool;
    
    // 对象大小监控服务
    private ObjectSizeMonitor objectSizeMonitor;
    
    // 上一次GC前的对象大小分布
    private long[] lastGcObjectDistribution;
    
    /**
     * 构造方法
     */
    public JmxMonitor() {
        this.memoryStatsHistory = new CopyOnWriteArrayList<>();
        this.gcStatsHistory = new CopyOnWriteArrayList<>();
        this.lastGcObjectDistribution = new long[6]; // [小对象数, 小对象大小, 中对象数, 中对象大小, 大对象数, 大对象大小]
    }
    
    /**
     * 设置对象大小监控服务
     */
    public void setObjectSizeMonitor(ObjectSizeMonitor objectSizeMonitor) {
        this.objectSizeMonitor = objectSizeMonitor;
    }
    
    /**
     * 启动JMX监控
     */
    public void start() {
        try {
            // 获取MBean服务器
            mbeanServer = ManagementFactory.getPlatformMBeanServer();
            
            // 获取内存MXBean
            memoryMXBean = ManagementFactory.getMemoryMXBean();
            
            // 获取内存池MXBean列表
            memoryPoolMXBeans = ManagementFactory.getMemoryPoolMXBeans();
            
            // 获取GC MXBean列表
            garbageCollectorMXBeans = ManagementFactory.getGarbageCollectorMXBeans();
            
            // 识别年轻代和老年代内存池
            identifyMemoryPools();
            
            // 注册GC通知监听器
            registerGcNotificationListener();
            
            // 初始化lastGcObjectDistribution数组，捕获初始对象分布
            if (objectSizeMonitor != null) {
                lastGcObjectDistribution = objectSizeMonitor.getCurrentObjectSizeDistribution();
                logger.info("Initialized lastGcObjectDistribution with current object distribution");
            } else {
                logger.info("ObjectSizeMonitor not set, lastGcObjectDistribution will be initialized with zeros");
            }
            
            logger.info("JMX Monitor started successfully");
        } catch (Exception e) {
            logger.error("Failed to start JMX Monitor: {}", e.getMessage(), e);
        }
    }
    
    /**
     * 停止JMX监控
     */
    public void stop() {
        try {
            // 注销GC通知监听器
            unregisterGcNotificationListener();
            logger.info("JMX Monitor stopped");
        } catch (Exception e) {
            logger.error("Failed to stop JMX Monitor: {}", e.getMessage(), e);
        }
    }
    
    /**
     * 识别年轻代和老年代内存池
     */
    private void identifyMemoryPools() {
        for (MemoryPoolMXBean pool : memoryPoolMXBeans) {
            String poolName = pool.getName().toLowerCase();
            MemoryType type = pool.getType();
            
            // 年轻代内存池识别
            if (type == MemoryType.HEAP && 
                (poolName.contains("eden") || 
                 poolName.contains("survivor") || 
                 poolName.contains("young") ||
                 poolName.contains("s0") || poolName.contains("s1"))) {
                youngMemoryPool = pool;
            }
            // 老年代内存池识别
            else if (type == MemoryType.HEAP && 
                     (poolName.contains("old") || 
                      poolName.contains("tenured") || 
                      poolName.contains("generation 1"))) {
                oldMemoryPool = pool;
            }
        }
        
        logger.info("Identified memory pools - Young: {}, Old: {}", 
                youngMemoryPool != null ? youngMemoryPool.getName() : "N/A",
                oldMemoryPool != null ? oldMemoryPool.getName() : "N/A");
    }
    
    /**
     * 注册GC通知监听器
     */
    private void registerGcNotificationListener() {
        try {
            // GC通知类型常量
            final String GC_NOTIFICATION_TYPE = "com.sun.management.gc.notification";
            
            // 创建GC通知监听器
            gcNotificationListener = new NotificationListener() {
                @Override
                public void handleNotification(Notification notification, Object handback) {
                    if (notification.getType().equals(GC_NOTIFICATION_TYPE)) {
                        handleGcNotification(notification);
                    }
                }
            };
            
            // 创建通知过滤器
            gcNotificationFilter = new NotificationFilter() {
                @Override
                public boolean isNotificationEnabled(Notification notification) {
                    return notification.getType().equals(GC_NOTIFICATION_TYPE);
                }
            };
            
            // 为每个GC MXBean注册通知监听器
            for (GarbageCollectorMXBean gcMxBean : garbageCollectorMXBeans) {
                ObjectName gcObjectName = getObjectName(gcMxBean);
                if (gcObjectName != null) {
                    mbeanServer.addNotificationListener(gcObjectName, gcNotificationListener, gcNotificationFilter, null);
                    logger.info("Registered GC notification listener for: {}", gcMxBean.getName());
                }
            }
        } catch (Exception e) {
            logger.error("Failed to register GC notification listener: {}", e.getMessage(), e);
        }
    }
    
    /**
     * 注销GC通知监听器
     */
    private void unregisterGcNotificationListener() {
        try {
            if (gcNotificationListener != null) {
                for (GarbageCollectorMXBean gcMxBean : garbageCollectorMXBeans) {
                    ObjectName gcObjectName = getObjectName(gcMxBean);
                    if (gcObjectName != null) {
                        mbeanServer.removeNotificationListener(gcObjectName, gcNotificationListener);
                    }
                }
                gcNotificationListener = null;
                logger.info("Unregistered GC notification listeners");
            }
        } catch (Exception e) {
            logger.error("Failed to unregister GC notification listener: {}", e.getMessage(), e);
        }
    }
    
    /**
     * 获取MXBean的ObjectName
     */
    private ObjectName getObjectName(GarbageCollectorMXBean gcMxBean) {
        try {
            // 直接使用MXBean的名称，不进行额外转义，让JMX API自行处理
            // 对于G1 GC等特殊名称，使用标准的MXBean获取方式
            return gcMxBean.getObjectName();
        } catch (Exception e) {
            try {
                // 备选方案：使用安全的ObjectName创建方式
                String gcName = gcMxBean.getName().replace(":", "_");
                return new ObjectName(ManagementFactory.GARBAGE_COLLECTOR_MXBEAN_DOMAIN_TYPE + ":name=" + gcName);
            } catch (MalformedObjectNameException ex) {
                logger.error("Failed to create ObjectName for GC MXBean '{}': {}", gcMxBean.getName(), ex.getMessage());
                return null;
            }
        }
    }
    
    /**
     * 处理GC通知
     */
    private void handleGcNotification(Notification notification) {
        try {
            // 从通知中获取GC信息
            CompositeData userData = (CompositeData) notification.getUserData();
            
            // 提取GC相关信息，使用安全的方式访问可能不存在的键
            String gcAction = safeGetString(userData, "gcAction");
            String gcCause = safeGetString(userData, "gcCause");
            
            // 使用当前时间作为默认值，处理不同JVM实现的GC通知格式差异
            long startTime = safeGetLong(userData, "gcStartTime", System.currentTimeMillis());
            long endTime = safeGetLong(userData, "gcEndTime", System.currentTimeMillis());
            
            // 获取内存池名称，使用安全访问
            String memoryPoolNames = safeGetString(userData, "memoryPoolNames");
            
            // 计算GC持续时间
            long duration = Math.max(0, endTime - startTime);
            
            // 收集GC前后的内存使用情况
            collectGcStats(gcAction, gcCause, startTime, duration, memoryPoolNames);
        } catch (Exception e) {
            logger.error("Failed to handle GC notification: {}", e.getMessage(), e);
        }
    }
    
    /**
     * 安全获取CompositeData中的String值
     */
    private String safeGetString(CompositeData data, String key) {
        try {
            return data != null && data.containsKey(key) ? (String) data.get(key) : "Unknown";
        } catch (Exception e) {
            return "Unknown";
        }
    }
    
    /**
     * 安全获取CompositeData中的long值
     */
    private long safeGetLong(CompositeData data, String key, long defaultValue) {
        try {
            return data != null && data.containsKey(key) ? (Long) data.get(key) : defaultValue;
        } catch (Exception e) {
            return defaultValue;
        }
    }
    
    /**
     * 收集GC统计数据
     */
    private void collectGcStats(String gcAction, String gcCause, long startTime, long duration, String memoryPoolNames) {
        try {
            GcStats gcStats = new GcStats();
            
            // 确定GC类型
            GcStats.GcType gcType = determineGcType(gcAction, gcCause);
            gcStats.setGcType(gcType);
            
            // 设置基本信息
            gcStats.setStartTime(startTime);
            gcStats.setDuration(duration);
            gcStats.setGcCause(gcCause);
            gcStats.setMemoryPoolNames(memoryPoolNames);
            
            // 获取当前内存使用情况（GC后的状态，因为通知是在GC完成后发送的）
            MemoryUsage heapAfter = memoryMXBean.getHeapMemoryUsage();
            long youngAfter = 0, oldAfter = 0;
            if (youngMemoryPool != null) {
                youngAfter = youngMemoryPool.getUsage().getUsed();
            }
            if (oldMemoryPool != null) {
                oldAfter = oldMemoryPool.getUsage().getUsed();
            }
            
            // 计算回收的内存量（使用峰值内存减去当前内存）
            // 这是一个更准确的估算方法，因为峰值通常出现在GC前
            long youngMemoryRecycled = calculatePoolMemoryRecycled(youngMemoryPool);
            long oldMemoryRecycled = calculatePoolMemoryRecycled(oldMemoryPool);
            
            // 计算堆内存回收总量
            long heapMemoryRecycled = youngMemoryRecycled + oldMemoryRecycled;
            
            // 计算GC前的内存使用情况（基于GC后的数据和回收量）
            long heapBefore = heapAfter.getUsed() + heapMemoryRecycled;
            long youngBefore = youngAfter + youngMemoryRecycled;
            long oldBefore = oldAfter + oldMemoryRecycled;
            
            // 设置内存回收数据
            gcStats.setHeapMemoryBefore(heapBefore);
            gcStats.setHeapMemoryAfter(heapAfter.getUsed());
            gcStats.setHeapMemoryRecycled(heapMemoryRecycled);
            
            gcStats.setYoungMemoryBefore(youngBefore);
            gcStats.setYoungMemoryAfter(youngAfter);
            gcStats.setYoungMemoryRecycled(youngMemoryRecycled);
            
            gcStats.setOldMemoryBefore(oldBefore);
            gcStats.setOldMemoryAfter(oldAfter);
            gcStats.setOldMemoryRecycled(oldMemoryRecycled);
            
            // 计算回收的对象大小分布
            calculateRecycledObjectStats(gcStats);
            
            // 添加到历史记录
            addGcStatsToHistory(gcStats);
            
            // 记录GC统计信息
            logGcStats(gcStats);
        } catch (Exception e) {
            logger.error("Failed to collect GC stats: {}", e.getMessage(), e);
        }
    }
    
    /**
     * 计算回收的对象大小统计
     */
    private void calculateRecycledObjectStats(GcStats gcStats) {
        if (objectSizeMonitor == null) {
            return;
        }
        
        // 获取GC后的对象大小分布
        long[] currentDistribution = objectSizeMonitor.getCurrentObjectSizeDistribution();
        
        // 计算回收的对象数量和大小（基于GC前后的差值）
        // 注意：这是一个近似计算，实际回收的对象可能与这个差值不同
        long smallObjectsRecycledCount = Math.max(0, lastGcObjectDistribution[0] - currentDistribution[0]);
        long smallObjectsRecycledSize = Math.max(0, lastGcObjectDistribution[1] - currentDistribution[1]);
        
        long mediumObjectsRecycledCount = Math.max(0, lastGcObjectDistribution[2] - currentDistribution[2]);
        long mediumObjectsRecycledSize = Math.max(0, lastGcObjectDistribution[3] - currentDistribution[3]);
        
        long largeObjectsRecycledCount = Math.max(0, lastGcObjectDistribution[4] - currentDistribution[4]);
        long largeObjectsRecycledSize = Math.max(0, lastGcObjectDistribution[5] - currentDistribution[5]);
        
        // 更新GC统计数据
        gcStats.setSmallObjectsRecycledCount(smallObjectsRecycledCount);
        gcStats.setSmallObjectsRecycledSize(smallObjectsRecycledSize);
        
        gcStats.setMediumObjectsRecycledCount(mediumObjectsRecycledCount);
        gcStats.setMediumObjectsRecycledSize(mediumObjectsRecycledSize);
        
        gcStats.setLargeObjectsRecycledCount(largeObjectsRecycledCount);
        gcStats.setLargeObjectsRecycledSize(largeObjectsRecycledSize);
        
        // 更新上一次GC前的对象分布
        System.arraycopy(currentDistribution, 0, lastGcObjectDistribution, 0, currentDistribution.length);
    }
    
    // 删除未使用的方法
    // 注释掉但保留代码，以备将来需要时恢复
    /*
    /**
     * 计算堆内存回收总量
     */
    /* 注释掉未使用的方法
    private long calculateHeapMemoryRecycled() {
        // 使用GC前后的堆内存差值计算回收量
        // 注意：这是一个近似值，因为GC可能会释放不同内存区域的内存
        MemoryUsage heapAfter = memoryMXBean.getHeapMemoryUsage();
        
        // 计算回收的内存量（基于GC前后的差值）
        // 这里假设我们在GC前保存了堆内存使用情况
        // 但由于我们没有实际保存，所以使用一个替代方案
        // 实际应用中，可以通过定期记录堆内存使用情况来计算更准确的回收量
        return 0;
    }
    */
    
    /**
     * 计算特定内存池的回收内存量
     */
    private long calculatePoolMemoryRecycled(MemoryPoolMXBean memoryPool) {
        if (memoryPool == null) {
            return 0;
        }
        
        // 使用内存池的峰值和当前使用量计算回收量
        // 这是一个更准确的估算方法，因为峰值通常出现在GC前
        MemoryUsage usage = memoryPool.getUsage();
        MemoryUsage peakUsage = memoryPool.getPeakUsage();
        
        long used = usage.getUsed();
        long peakUsed = peakUsage.getUsed();
        
        // 计算回收的内存量（峰值减去当前使用量）
        long recycled = Math.max(0, peakUsed - used);
        
        // 重置峰值使用量，以便下一次计算
        memoryPool.resetPeakUsage();
        
        return recycled;
    }
    
    /**
     * 确定GC类型
     */
    private GcStats.GcType determineGcType(String gcAction, String gcCause) {
        if (gcAction.contains("minor") || gcAction.contains("Young") || gcAction.contains("young")) {
            return GcStats.GcType.YOUNG_GC;
        } else if (gcAction.contains("major") || gcAction.contains("Old") || gcAction.contains("old")) {
            return GcStats.GcType.OLD_GC;
        } else if (gcAction.contains("full") || gcAction.contains("Full") || gcCause.contains("full") || gcCause.contains("Full")) {
            return GcStats.GcType.FULL_GC;
        } else {
            return GcStats.GcType.UNKNOWN;
        }
    }
    
    /**
     * 记录GC统计信息
     */
    private void logGcStats(GcStats gcStats) {
        // 使用结构化的单条日志记录GC信息，减少日志行数并提高可读性
        String gcInfo = String.format("[GC] %s | Duration: %dms | Cause: %s | Heap: %s (Before: %s, After: %s) | " +
                        "Young: %s, Old: %s | Objects: Small(%d/%s), Medium(%d/%s), Large(%d/%s)",
                gcStats.getGcType(), 
                gcStats.getDuration(), 
                gcStats.getGcCause(),
                formatBytes(gcStats.getHeapMemoryRecycled()),
                formatBytes(gcStats.getHeapMemoryBefore()),
                formatBytes(gcStats.getHeapMemoryAfter()),
                formatBytes(gcStats.getYoungMemoryRecycled()),
                formatBytes(gcStats.getOldMemoryRecycled()),
                gcStats.getSmallObjectsRecycledCount(), 
                formatBytes(gcStats.getSmallObjectsRecycledSize()),
                gcStats.getMediumObjectsRecycledCount(), 
                formatBytes(gcStats.getMediumObjectsRecycledSize()),
                gcStats.getLargeObjectsRecycledCount(), 
                formatBytes(gcStats.getLargeObjectsRecycledSize()));
        
        logger.info(gcInfo);
    }
    
    /**
     * 收集内存统计数据
     */
    public void collectMemoryStats() {
        try {
            MemoryStats stats = new MemoryStats();
            stats.setTimestamp(System.currentTimeMillis());
            
            // 获取堆内存信息
            MemoryUsage heapMemoryUsage = memoryMXBean.getHeapMemoryUsage();
            stats.setHeapMemoryUsed(heapMemoryUsage.getUsed());
            stats.setHeapMemoryMax(heapMemoryUsage.getMax());
            stats.setHeapMemoryCommitted(heapMemoryUsage.getCommitted());
            
            // 获取非堆内存信息
            MemoryUsage nonHeapMemoryUsage = memoryMXBean.getNonHeapMemoryUsage();
            stats.setNonHeapMemoryUsed(nonHeapMemoryUsage.getUsed());
            stats.setNonHeapMemoryMax(nonHeapMemoryUsage.getMax());
            stats.setNonHeapMemoryCommitted(nonHeapMemoryUsage.getCommitted());
            
            // 遍历内存池，获取更详细的内存信息
            for (MemoryPoolMXBean pool : memoryPoolMXBeans) {
                String poolName = pool.getName();
                MemoryUsage usage = pool.getUsage();
                
                if (usage == null) {
                    continue;
                }
                
                // 元空间
                if (poolName.contains("Metaspace")) {
                    stats.setMetaspaceUsed(usage.getUsed());
                    stats.setMetaspaceMax(usage.getMax());
                    stats.setMetaspaceCommitted(usage.getCommitted());
                }
                // 代码缓存
                else if (poolName.contains("Code Cache")) {
                    stats.setCodeCacheUsed(usage.getUsed());
                    stats.setCodeCacheMax(usage.getMax());
                    stats.setCodeCacheCommitted(usage.getCommitted());
                }
            }
            
            // 获取直接内存使用情况
            long directMemoryUsed = getDirectMemoryUsed();
            stats.setDirectMemoryUsed(directMemoryUsed);
            
            // 更新最新统计数据
            latestMemoryStats = stats;
            
            // 添加到历史记录
            addToHistory(stats);
            
            logger.debug("Collected memory stats: {}", stats);
        } catch (Exception e) {
            logger.error("Failed to collect memory stats: {}", e.getMessage(), e);
        }
    }
    
    /**
     * 获取直接内存使用情况
     */
    private long getDirectMemoryUsed() {
        try {
            // 使用JMX获取直接内存使用情况
            ObjectName directMemoryMBeanName = new ObjectName("java.nio:type=BufferPool,name=direct");
            Object memoryUsed = mbeanServer.getAttribute(directMemoryMBeanName, "MemoryUsed");
            
            if (memoryUsed instanceof Number) {
                return ((Number) memoryUsed).longValue();
            }
            
            logger.debug("Unexpected type for MemoryUsed attribute: {}", memoryUsed.getClass().getName());
            return 0;
        } catch (Exception e) {
            logger.debug("Failed to get direct memory usage: {}", e.getMessage());
            return 0;
        }
    }
    
    /**
     * 将内存统计数据添加到历史记录
     */
    private void addToHistory(MemoryStats stats) {
        memoryStatsLock.writeLock().lock();
        try {
            memoryStatsHistory.add(stats);
            
            // 保持历史记录大小不超过最大值
            if (memoryStatsHistory.size() > MAX_HISTORY_SIZE) {
                memoryStatsHistory.remove(0);
            }
        } finally {
            memoryStatsLock.writeLock().unlock();
        }
    }
    
    /**
     * 将GC统计数据添加到历史记录
     */
    private void addGcStatsToHistory(GcStats gcStats) {
        gcStatsLock.writeLock().lock();
        try {
            gcStatsHistory.add(gcStats);
            
            // 保持历史记录大小不超过最大值
            if (gcStatsHistory.size() > MAX_HISTORY_SIZE) {
                gcStatsHistory.remove(0);
            }
        } finally {
            gcStatsLock.writeLock().unlock();
        }
    }
    
    /**
     * 获取最新的内存统计数据
     */
    public MemoryStats getLatestMemoryStats() {
        return latestMemoryStats;
    }
    
    /**
     * 获取内存统计历史记录
     */
    public List<MemoryStats> getMemoryStatsHistory() {
        memoryStatsLock.readLock().lock();
        try {
            return new ArrayList<>(memoryStatsHistory);
        } finally {
            memoryStatsLock.readLock().unlock();
        }
    }
    
    /**
     * 获取GC统计历史记录
     */
    public List<GcStats> getGcStatsHistory() {
        gcStatsLock.readLock().lock();
        try {
            return new ArrayList<>(gcStatsHistory);
        } finally {
            gcStatsLock.readLock().unlock();
        }
    }
    
    /**
     * 计算非堆内存占总内存的比例
     */
    public double getNonHeapMemoryRatio() {
        if (latestMemoryStats == null) {
            return 0.0;
        }
        
        long totalMemory = latestMemoryStats.getHeapMemoryUsed() + latestMemoryStats.getNonHeapMemoryUsed();
        if (totalMemory == 0) {
            return 0.0;
        }
        
        return (double) latestMemoryStats.getNonHeapMemoryUsed() / totalMemory;
    }
    
    /**
     * 计算元空间占非堆内存的比例
     */
    public double getMetaspaceRatio() {
        if (latestMemoryStats == null || latestMemoryStats.getNonHeapMemoryUsed() == 0) {
            return 0.0;
        }
        
        return (double) latestMemoryStats.getMetaspaceUsed() / latestMemoryStats.getNonHeapMemoryUsed();
    }
    
    /**
     * 计算直接内存占非堆内存的比例
     */
    public double getDirectMemoryRatio() {
        if (latestMemoryStats == null || latestMemoryStats.getNonHeapMemoryUsed() == 0) {
            return 0.0;
        }
        
        return (double) latestMemoryStats.getDirectMemoryUsed() / latestMemoryStats.getNonHeapMemoryUsed();
    }
    
    /**
     * 格式化字节数为可读格式
     */
    private String formatBytes(long bytes) {
        if (bytes < 1024) {
            return bytes + "B";
        } else if (bytes < 1024 * 1024) {
            return String.format("%.1fKB", bytes / 1024.0);
        } else if (bytes < 1024 * 1024 * 1024) {
            return String.format("%.1fMB", bytes / (1024.0 * 1024.0));
        } else {
            return String.format("%.1fGB", bytes / (1024.0 * 1024.0 * 1024.0));
        }
    }
}